#! py -3
# coding:utf-8

import numpy as np
import cv2 as cv
import os
from os import path
import torch.utils.data
import torchvision


def getFilepathInDir(dirstr: str):
    assert type(dirstr) is str
    if dirstr is None:
        return []
    if not path.exists(dirstr):
        return []
    if not path.isdir(dirstr):
        return []
    filter = [".PNG", ".JPG", ".BMP"]
    ret = []
    filelist = os.listdir(dirstr)
    for s in filelist:
        ext = path.splitext(s)[1]
        ext = ext.upper()
        if ext in filter:
            ret.append(s)
            # print(s)
    ret.sort()
    for i, s in enumerate(ret):
        ret[i] = path.join(dirstr, s).replace("\\", "/")
    return ret


def compareIsHasSameFilename(filelist1: str, filelist2: str):
    assert type(filelist1) is list
    assert type(filelist2) is list
    if len(filelist1) != len(filelist2):
        return False
    cnt = len(filelist1)
    isequal = True
    for i in range(cnt):
        n1 = path.split(filelist1[i])[1]
        n2 = path.split(filelist2[i])[1]
        # print(n1, n2)
        isequal = isequal and (n1 == n2)
        if not isequal:
            break
    return isequal


# 直接读取图像文件名列表 作为数据集
class ImgListDataSet(torch.utils.data.Dataset):
    def __init__(self, imagelist=[], reqsize=(512, 512), trans: torchvision.transforms.Compose = None):
        self.imagelist = imagelist
        self.reqsize = reqsize
        self.transforms = trans

    def __len__(self):
        return len(self.imagelist)

    def __getitem__(self, index):
        count = len(self)
        image = None
        if index < count or count < 1:
            image = cv.imread(self.imagelist[index], cv.IMREAD_GRAYSCALE)

        if image is None:
            image = np.zeros((512, 512, 1), dtype=np.uint8)

        if not self.reqsize is None:
            image = cv.resize(image, self.reqsize)

        if not self.transforms is None:
            image = self.transforms(image)
        return image


class OCT_DataSet(torch.utils.data.Dataset):
    def __init__(self,
                 train_inputFilelist=[],
                 train_outputFilelist=[],
                 validate_inputFilelist=[],
                 validate_outputFilelist=[],
                 imgsize=None,
                 trainmode=True,
                 xtransforms: torchvision.transforms.Compose = None,
                 # =torchvision.transforms.Compose([torchvision.transforms.ToTensor()])
                 ytransforms: torchvision.transforms.Compose = None
                 ):
        self.train_inputFilelist = train_inputFilelist
        self.train_outputFilelist = train_outputFilelist
        self.validate_inputFilelist = validate_inputFilelist
        self.validate_outputFilelist = validate_outputFilelist
        self.imgsize = imgsize
        self.trainmode = trainmode
        self.xtransforms = xtransforms
        self.ytransforms = ytransforms
        return

    def __len__(self):
        cnt = 0
        if self.trainmode:
            cnt = min(len(self.train_inputFilelist), len(self.train_outputFilelist))
            return cnt
        else:
            cnt = min(len(self.validate_inputFilelist), len(self.validate_outputFilelist))
            return cnt

    def __getitem__(self, index):
        cnt = len(self)
        if index >= cnt:
            raise StopIteration
        xlist = self.train_inputFilelist
        ylist = self.train_outputFilelist
        if not self.trainmode:
            xlist = self.validate_inputFilelist
            ylist = self.validate_outputFilelist

        size = self.imgsize
        if type(size) is not tuple:
            size = (512, 512)
        if len(size) < 2:
            size = (512, 512)

        while True:
            ximg = None
            yimg = None
            ximg = cv.imread(xlist[index], cv.IMREAD_GRAYSCALE)
            yimg = cv.imread(ylist[index], cv.IMREAD_GRAYSCALE)
            if (ximg is None) or (yimg is None):
                index = index + 1
                if index >= cnt:
                    raise StopIteration
            else:
                break

        ximg: np.ndarray
        yimg: np.ndarray
        if ximg.dtype != np.uint8:
            ximg = ximg.astype(np.uint8)
        if yimg.dtype != np.uint8:
            yimg = yimg.astype(np.uint8)
        if not (self.imgsize is None):
            ximg = cv.resize(ximg, size)
            yimg = cv.resize(ximg, size)

        if ximg.shape != yimg.shape:
            ximg = cv.resize(ximg, yimg.shape)

        ximg = np.expand_dims(ximg, 2)
        yimg = np.expand_dims(yimg, 2)
        if not (self.xtransforms is None):
            ximg = self.xtransforms(ximg)
        if not (self.ytransforms is None):
            yimg = self.ytransforms(yimg)

        return ximg, yimg

# inputlist = getFilepathInDir("../oct_img/input")
# # print(inputlist)
# outputlist = getFilepathInDir("../oct_img/output")
#
# input_vallist = getFilepathInDir("../oct_img/input_val")
# output_vallist = getFilepathInDir("../oct_img/output_val")
#
# print(compareIsHasSameFilename(inputlist, outputlist))
# print(compareIsHasSameFilename(input_vallist, output_vallist))

# dataset = ImgListDataSet(inputlist)
# cv.imshow("image", dataset[0])
# cv.waitKey(0)

# dataset = OCT_DataSet(inputlist, outputlist, input_vallist, output_vallist, None, False, None)
# for x, y in dataset:
#     cv.imshow("x", x)
#     cv.imshow("y", y)
#     cv.waitKey(33)
#
# cv.destroyAllWindows()

# # tensor 可以用permute 掉换维度
# # numpy 可以用transpose 掉换维度

# dataset = OCT_DataSet(inputlist, outputlist, input_vallist, output_vallist)
# dataloader = torch.utils.data.DataLoader(dataset, 10, shuffle=True)
# for x, y in dataloader:
#     ximg = x[0].numpy()
#     yimg = y[0].numpy()
#     cv.imshow("x", ximg)
#     cv.imshow("y", yimg)
#     cv.waitKey(33)
#     # print(x, y)
